from collections import deque
from collections.abc import Iterator
from threading import Condition, RLock, Thread, current_thread
from typing import Any, TypeVar

from zkl_aiutils_datasets.basics import Dataset, DatasetIterator, DatasetWriter, resume
from .state_safe import StateSafeTransformedIterator

AnySample = TypeVar("AnySample")
AnyState = TypeVar('AnyState')


class ThreadedBufferedWriter(DatasetWriter[AnySample]):
    def __init__(self, writer: DatasetWriter[AnySample], buffer_size: int = 1024):
        self._writer = writer
        self._buffer = deque(maxlen=buffer_size)
        self._condition = Condition(RLock())
        self._closed_ref = [False]
        self._exception_ref = [None]
        self._thread = Thread(
            target=self._thread_func, daemon=True,
            args=(self._writer, self._buffer, self._condition, self._closed_ref, self._exception_ref))
        self._thread.start()

    @classmethod
    def _thread_func(cls,
        writer: DatasetWriter[AnySample],
        buffer: deque[AnySample],
        condition: Condition,
        closed_ref: list[bool],
        exception_ref: list[Exception],
    ):
        while True:
            with condition:
                try:
                    condition.wait_for(lambda: any((
                        len(buffer) > 0,
                        closed_ref[0])))
                    if len(buffer) > 0:
                        sample = buffer.popleft()
                        writer.write(sample)
                        condition.notify_all()
                    elif closed_ref[0]:
                        break
                except Exception as exception:
                    exception_ref[0] = exception
                    condition.notify_all()
                    break

    def write(self, sample: AnySample):
        with self._condition:
            self._condition.wait_for(lambda: any((
                self._exception_ref[0] is not None,
                self._closed_ref[0],
                len(self._buffer) < self._buffer.maxlen)))
            if self._exception_ref[0]:
                raise self._exception_ref[0]
            if self._closed_ref[0]:
                raise RuntimeError(f"Cannot write when the writer is closed.")
            self._buffer.append(sample)
            self._condition.notify_all()

    def flush(self):
        with self._condition:
            self._condition.wait_for(lambda: any((
                self._exception_ref[0] is not None,
                self._closed_ref[0],
                len(self._buffer) == 0)))
            if self._exception_ref[0]:
                raise self._exception_ref[0]
            if self._closed_ref[0]:
                raise RuntimeError(f"Cannot flush when the writer is closed.")
            self._writer.flush()

    def close(self):
        with self._condition:
            self._closed_ref[0] = True
            self._condition.notify_all()
            self._condition.wait_for(lambda: any((
                self._exception_ref[0] is not None,
                len(self._buffer) == 0)))
            if self._exception_ref[0]:
                raise self._exception_ref[0]
            self._writer.close()
        if current_thread() is not self._thread:
            self._thread.join()


class ThreadedBufferedIterator(DatasetIterator[AnySample]):
    def __init__(self, iterator: Iterator[AnySample], buffer_size: int = 1024):
        self._iterator = iterator
        self._buffer = deque(maxlen=buffer_size)
        self._condition = Condition(RLock())
        self._closed_ref = [False]
        self._exception_ref = [None]
        self._thread = Thread(
            target=self._thread_func, daemon=True,
            args=(self._iterator, self._buffer, self._condition, self._closed_ref, self._exception_ref))
        self._thread.start()

    @classmethod
    def _thread_func(cls,
        iterator: Iterator[AnySample],
        buffer: deque,
        condition: Condition,
        closed_ref: list[bool],
        exception_ref: list[Exception],
    ):
        while True:
            try:
                sample = next(iterator)
            except Exception as exception:
                with condition:
                    exception_ref[0] = exception
                    condition.notify_all()
                break

            with condition:
                condition.wait_for(lambda: any((
                    closed_ref[0],
                    len(buffer) < buffer.maxlen)))
                if closed_ref[0]:
                    break
                buffer.append(sample)
                condition.notify_all()

    def __next__(self) -> AnySample:
        with self._condition:
            self._condition.wait_for(lambda: any((
                len(self._buffer) > 0,
                self._exception_ref[0],
                self._closed_ref[0])))
            if len(self._buffer) > 0:
                sample = self._buffer.popleft()
                self._condition.notify_all()
                return sample
            if self._exception_ref[0]:
                raise self._exception_ref[0]
            raise RuntimeError(f"{self} is not usable after calling method `__del__()`.")

    def __del__(self):
        with self._condition:
            self._closed_ref[0] = True
            self._condition.notify_all()
        if current_thread() is not self._thread:
            self._thread.join()


class ThreadedBufferedDataset(Dataset[AnySample]):
    def __init__(self, dataset: Dataset[AnySample], buffer_size: int = 1024):
        self._dataset = dataset
        self._buffer_size = buffer_size

    @property
    def dataset(self) -> Dataset[AnySample]:
        return self._dataset

    @property
    def buffer_size(self) -> int:
        return self._buffer_size

    def __len__(self) -> int:
        return len(self._dataset)

    def __getitem__(self, index: int, /) -> AnySample:
        return self._dataset[index]

    def __iter__(self) -> DatasetIterator[AnySample]:
        iterator = iter(self._dataset)
        return ThreadedBufferedIterator(iterator, self._buffer_size)


class StateSafeThreadedBufferedIterator(StateSafeTransformedIterator[AnySample, AnyState]):
    def __init__(self, iterator: Iterator[AnySample], buffer_size: int = 1024):
        super().__init__(iterator, lambda it: ThreadedBufferedIterator(it, buffer_size))


class StateSafeThreadedBufferedDataset(ThreadedBufferedDataset[AnySample]):
    def __iter__(self) -> DatasetIterator[AnySample]:
        return StateSafeThreadedBufferedIterator(iter(self._dataset), self._buffer_size)

    def __resume__(self, state: Any | None = None, /) -> DatasetIterator[AnySample]:
        return StateSafeThreadedBufferedIterator(resume(self._dataset, state), self._buffer_size)
